import torch
import torch.nn as nn
import math
from .base_model import BaseModel  # 修正导入：从同级目录导入BaseModel
from . import networks
from torchmetrics.functional import structural_similarity_index_measure as ssim


# ------------------------------
# 1. ViT注意力模块（仅负责注意力计算，继承nn.Module）
# ------------------------------
class ViTBlock(nn.Module):
    """改进的ViT块：仅保留关键位置编码+多头注意力"""

    def __init__(self, dim, num_heads=4, mlp_ratio=4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(dim)  # 注意力前归一化
        self.attn = nn.MultiheadAttention(dim, num_heads, batch_first=True)  # 多头注意力
        self.norm2 = nn.LayerNorm(dim)  # MLP前归一化
        # MLP层（GELU激活更适合ViT）
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Linear(int(dim * mlp_ratio), dim)
        )

    def get_positional_encoding(self, seq_len, dim, device):
        """生成位置编码（确保与输入设备一致，避免CPU/GPU不匹配）"""
        position = torch.arange(seq_len, device=device).unsqueeze(1)  # [seq_len, 1]
        div_term = torch.exp(torch.arange(0, dim, 2, device=device) * (-math.log(10000.0) / dim))  # [dim//2]
        pos_encoding = torch.zeros(seq_len, dim, device=device)  # [seq_len, dim]
        pos_encoding[:, 0::2] = torch.sin(position * div_term)  # 偶数位用正弦
        pos_encoding[:, 1::2] = torch.cos(position * div_term)  # 奇数位用余弦
        return pos_encoding.unsqueeze(0)  # [1, seq_len, dim]（适配batch维度）

    def forward(self, x):
        """x: [B, seq_len, dim]（B=批量，seq_len=序列长度，dim=特征维度）"""
        seq_len, dim = x.shape[1], x.shape[2]
        # 1. 添加位置编码
        pos_encoding = self.get_positional_encoding(seq_len, dim, x.device)
        x = x + pos_encoding  # 残差连接：特征+位置编码

        # 2. 多头注意力（残差连接）
        x_res = x  # 保存残差
        x = self.norm1(x)
        x, _ = self.attn(x, x, x)  # 自注意力
        x = x + x_res  # 残差更新

        # 3. MLP（残差连接）
        x_res = x  # 保存残差
        x = self.norm2(x)
        x = self.mlp(x)
        x = x + x_res  # 残差更新

        return x  # [B, seq_len, dim]


class ViTAttention(nn.Module):
    """ViT注意力封装：适配U-Net的2D特征图输入（.patch→ViT→unpatch）"""

    def __init__(self, in_channels, patch_size=4, num_heads=4):
        super().__init__()
        self.patch_size = patch_size
        self.dim = in_channels  # 特征维度=输入通道数（确保与U-Net瓶颈层通道一致）

        # 1. 特征图→Patch序列（Conv2d实现：步长=patch_size，输出通道=dim）
        self.patch_embed = nn.Conv2d(
            in_channels=in_channels,
            out_channels=self.dim,
            kernel_size=patch_size,
            stride=patch_size,
            padding=0
        )

        # 2. ViT块（处理Patch序列）
        self.vit_block = ViTBlock(dim=self.dim, num_heads=num_heads)

        # 3. Patch序列→特征图（转置卷积恢复尺寸）
        self.unpatch = nn.ConvTranspose2d(
            in_channels=self.dim,
            out_channels=in_channels,
            kernel_size=patch_size,
            stride=patch_size,
            padding=0
        )

    def forward(self, x):
        """x: [B, C, H, W]（U-Net瓶颈层输出，C=in_channels）"""
        B, C, H, W = x.shape

        # 1. 特征图→Patch序列：[B, C, H, W] → [B, dim, H/patch, W/patch] → [B, seq_len, dim]
        x_patch = self.patch_embed(x)  # [B, dim, H/patch, W/patch]
        seq_len = (H // self.patch_size) * (W // self.patch_size)  # 序列长度=Patch数量
        x_seq = x_patch.flatten(2).transpose(1, 2)  # [B, seq_len, dim]（batch_first=True）

        # 2. ViT注意力计算
        x_seq = self.vit_block(x_seq)  # [B, seq_len, dim]

        # 3. Patch序列→特征图：[B, seq_len, dim] → [B, dim, H/patch, W/patch] → [B, C, H, W]
        x_patch = x_seq.transpose(1, 2).view(B, self.dim, H // self.patch_size, W // self.patch_size)
        x = self.unpatch(x_patch)  # [B, C, H, W]（与输入尺寸完全一致）

        return x


# ------------------------------
# 2. 完整训练模型（继承BaseModel，类名匹配--model pix2pixwithvit）
# ------------------------------
class Pix2pixwithvitModel(BaseModel):
    """Pix2Pix + ViT注意力模型（可直接训练）"""

    @staticmethod
    def modify_commandline_options(parser, is_train=True):
        """添加模型专属参数，覆盖默认配置（与参考模型一致）"""
        # 基础配置：匹配Pix2Pix默认（batch norm、U-Net256、aligned数据集）
        parser.set_defaults(norm='batch', netG='unet_256', dataset_mode='aligned')
        if is_train:
            # 训练配置：无图像缓冲、Vanilla GAN、L1/SSIM损失权重
            parser.set_defaults(pool_size=0, gan_mode='vanilla')
            parser.add_argument('--lambda_L1', type=float, default=100.0, help='L1损失权重')
            parser.add_argument('--lambda_SSIM', type=float, default=0.1, help='SSIM损失权重（保守值）')
        return parser

    def __init__(self, opt):
        """初始化模型：定义生成器、判别器、损失函数、优化器"""
        BaseModel.__init__(self, opt)  # 调用BaseModel初始化（设备、模型名列表等）

        # 1. 定义损失名（训练时需打印/保存的损失）
        self.loss_names = ['G_GAN', 'G_L1', 'G_SSIM', 'D_real', 'D_fake']
        # 2. 定义可视化名（训练时需显示的图像）
        self.visual_names = ['real_A', 'fake_B', 'real_B']
        # 3. 定义模型名（训练时需保存/加载的模型）
        self.model_names = ['G', 'D'] if self.isTrain else ['G']  # 测试时仅需生成器

        # ------------------------------
        # 4. 定义生成器（U-Net + ViT注意力）
        # ------------------------------
        # 4.1 基础U-Net生成器（调用networks.define_G，与原生Pix2Pix一致）
        self.netG = networks.define_G(
            input_nc=opt.input_nc,  # 输入通道数（如RGB=3）
            output_nc=opt.output_nc,  # 输出通道数（如RGB=3）
            ngf=opt.ngf,  # 生成器基础通道数（默认64）
            netG=opt.netG,  # 生成器结构（默认unet_256）
            norm=opt.norm,  # 归一化方式（默认batch）
            use_dropout=not opt.no_dropout,  # 是否用dropout
            init_type=opt.init_type,  # 权重初始化方式（默认normal）
            init_gain=opt.init_gain,  # 初始化增益（默认0.02）
            gpu_ids=self.gpu_ids  # GPU设备列表
        )

        # 4.2 在U-Net瓶颈层插入ViT注意力（关键修改）
        # 注：unet_256的瓶颈层通道数=ngf*8（ngf=64时为512），需确保ViT输入通道匹配
        if hasattr(self.netG, 'bottleneck'):  # 确认U-Net有bottleneck属性
            original_bottleneck = self.netG.bottleneck  # 保存原瓶颈层
            # 替换为「原瓶颈层 + ViT注意力」（残差增强，避免破坏原有特征）
            self.netG.bottleneck = nn.Sequential(
                original_bottleneck,
                ViTAttention(in_channels=opt.ngf * 8, patch_size=4)  # in_channels=ngf*8（适配U-Net）
            )
            print("✅ ViT注意力已插入U-Net瓶颈层")

        # ------------------------------
        # 5. 定义判别器（仅训练时需要）
        # ------------------------------
        if self.isTrain:
            self.netD = networks.define_D(
                input_nc=opt.input_nc + opt.output_nc,  # 输入=A+B（拼接真实输入+生成输出）
                ndf=opt.ndf,  # 判别器基础通道数（默认64）
                netD=opt.netD,  # 判别器结构（默认basic/PatchGAN）
                n_layers_D=opt.n_layers_D,  # 判别器层数（默认3）
                norm=opt.norm,  # 归一化方式
                init_type=opt.init_type,
                init_gain=opt.init_gain,
                gpu_ids=self.gpu_ids
            )

            # ------------------------------
            # 6. 定义损失函数
            # ------------------------------
            self.criterionGAN = networks.GANLoss(opt.gan_mode).to(self.device)  # GAN损失（Vanilla）
            self.criterionL1 = torch.nn.L1Loss()  # L1损失（像素级匹配）
            self.ssim_weight = getattr(opt, 'lambda_SSIM', 0.1)  # SSIM损失权重（默认0.1）

            # ------------------------------
            # 7. 定义优化器
            # ------------------------------
            # 生成器优化器：学习率降低至0.5*opt.lr（避免ViT模块训练不稳定）
            self.optimizer_G = torch.optim.Adam(
                self.netG.parameters(),
                lr=opt.lr * 0.5,
                betas=(opt.beta1, 0.999)  # beta1=0.5（GAN常用配置）
            )
            # 判别器优化器：学习率=opt.lr（与原生Pix2Pix一致）
            self.optimizer_D = torch.optim.Adam(
                self.netD.parameters(),
                lr=opt.lr,
                betas=(opt.beta1, 0.999)
            )
            # 将优化器加入BaseModel的optimizers列表（自动管理训练步骤）
            self.optimizers.append(self.optimizer_G)
            self.optimizers.append(self.optimizer_D)  # 修正笔误：optimizers_D → optimizer_D

        print("✅ Pix2pixwithvitModel初始化完成")

    # ------------------------------
    # 8. 数据输入处理（与参考模型完全一致）
    # ------------------------------
    def set_input(self, input):
        """input: 数据加载器返回的字典（包含A、B、路径）"""
        AtoB = self.opt.direction == 'AtoB'  # 方向：A→B 或 B→A
        self.real_A = input['A' if AtoB else 'B'].to(self.device)  # 输入图像
        self.real_B = input['B' if AtoB else 'A'].to(self.device)  # 真实目标图像
        self.image_paths = input['A_paths' if AtoB else 'B_paths']  # 图像路径（用于日志）

    # ------------------------------
    # 9. 前向传播（生成器推理）
    # ------------------------------
    def forward(self):
        """生成fake_B：real_A → netG → fake_B"""
        self.fake_B = self.netG(self.real_A)  # [B, C, H, W]

    # ------------------------------
    # 10. 判别器反向传播（计算D损失）
    # ------------------------------
    def backward_D(self):
        """判别器目标：区分「real_A+real_B」和「real_A+fake_B」"""
        # 1. 处理fake_B（冻结生成器，避免梯度传到G）
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)  # 拼接：[B, C*2, H, W]
        pred_fake = self.netD(fake_AB.detach())  # detach()：不计算G的梯度
        self.loss_D_fake = self.criterionGAN(pred_fake, False)  # 判别器认为fake是"假"

        # 2. 处理real_B
        real_AB = torch.cat((self.real_A, self.real_B), 1)
        pred_real = self.netD(real_AB)
        self.loss_D_real = self.criterionGAN(pred_real, True)  # 判别器认为real是"真"

        # 3. 总判别器损失（平均fake和real损失）
        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5
        self.loss_D.backward()  # 反向传播：计算D的梯度

    # ------------------------------
    # 11. 生成器反向传播（计算G损失）
    # ------------------------------
    def backward_G(self):
        """生成器目标：1. 欺骗判别器；2. 生成图贴近真实图（L1+SSIM）"""
        # 1. GAN损失（欺骗判别器）
        fake_AB = torch.cat((self.real_A, self.fake_B), 1)
        pred_fake = self.netD(fake_AB)
        self.loss_G_GAN = self.criterionGAN(pred_fake, True)  # 让判别器认为fake是"真"

        # 2. L1损失（像素级匹配）
        self.loss_G_L1 = self.criterionL1(self.fake_B, self.real_B) * self.opt.lambda_L1

        # 3. SSIM损失（结构相似性，鲁棒计算避免报错）
        try:
            # SSIM要求输入尺寸≥11x11（避免小尺寸报错）
            if self.fake_B.shape[2] >= 11 and self.fake_B.shape[3] >= 11:
                # SSIM范围[0,1]，损失=1-SSIM（值越小结构越相似）
                self.loss_G_SSIM = (1 - ssim(self.fake_B, self.real_B, data_range=1.0)) * self.ssim_weight
            else:
                self.loss_G_SSIM = torch.tensor(0.0, device=self.device)  # 小尺寸时SSIM损失为0
        except Exception as e:
            print(f"⚠️ SSIM计算警告: {e}，临时将SSIM损失设为0")
            self.loss_G_SSIM = torch.tensor(0.0, device=self.device)

        # 4. 总生成器损失（GAN + L1 + SSIM）
        self.loss_G = self.loss_G_GAN + self.loss_G_L1 + self.loss_G_SSIM
        self.loss_G.backward()  # 反向传播：计算G的梯度

    # ------------------------------
    # 12. 优化步骤（单次训练迭代）
    # ------------------------------
    def optimize_parameters(self):
        """1. 前向传播生成fake_B；2. 优化判别器；3. 优化生成器"""
        # 1. 前向传播（生成fake_B）
        self.forward()

        # 2. 优化判别器（需开启D的梯度）
        self.set_requires_grad(self.netD, True)  # 允许D参数更新
        self.optimizer_D.zero_grad()  # 清空D梯度
        self.backward_D()  # 计算D损失+梯度
        self.optimizer_D.step()  # 更新D参数

        # 3. 优化生成器（冻结D，避免D干扰G的梯度）
        self.set_requires_grad(self.netD, False)  # 禁止D参数更新
        self.optimizer_G.zero_grad()  # 清空G梯度
        self.backward_G()  # 计算G损失+梯度
        self.optimizer_G.step()  # 更新G参数